# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
import os
import sys

from absl import app, flags
import pandas

from coin.support.proto_log.logic.util import (run_from_og_log_archive,
                                               resolve_trading_date_range_string)
from coin.proto.coin_order_gateway_pb2 import OrderEvent, OrderGatewayLog
from coin.support.proto_log.logic.pnl_calculator import PnlCalculatorFactory
from coin.support.proto_log.logic.stats_collector import OrderEventStatsCollector

FLAGS = flags.FLAGS


class AnalyzeStrategy(object):
  def __init__(self, exchange):
    self._exchange = exchange
    self._order_event_stats_collector = OrderEventStatsCollector()
    self._pnl_calculator = PnlCalculatorFactory(exchange, 0, 0).create_calculator()

  def on_log(self, timestamp, log):
    pb = OrderGatewayLog()
    pb.ParseFromString(log)
    if pb.type == OrderGatewayLog.ORDER_EVENT:
      event = pb.event
      self._order_event_stats_collector.update_by_order_event(event)
      if event.type == OrderEvent.ORDER_FILLED:
        self._pnl_calculator.update_by_fill(event)
    elif pb.type == OrderGatewayLog.SUBMIT_REQUEST:
      request = pb.submit_request
      self._order_event_stats_collector.update_by_submit_request(request, timestamp)
    elif pb.type == OrderGatewayLog.AMEND_REQUEST:
      request = pb.amend_request
      self._order_event_stats_collector.update_by_amend_request(request, timestamp)
    elif pb.type == OrderGatewayLog.CANCEL_REQUEST:
      request = pb.cancel_request
      self._order_event_stats_collector.update_by_cancel_request(request, timestamp)
    else:
      # Don't care, skip.
      pass

  def pnl_result(self):
    pnl_dict = {
        'realized_pnl': self._pnl_calculator.realized_pnl,
        'latest_position': self._pnl_calculator.latest_pos,
        'latest_position_price': self._pnl_calculator.latest_pos_price
    }
    return pnl_dict

  def fill_ratio_result(self):
    total_post_count = self._order_event_stats_collector.total_submit_count - self._order_event_stats_collector.total_submit_reject_count
    total_post_qty = self._order_event_stats_collector.total_submit_qty - self._order_event_stats_collector.total_submit_reject_qty
    total_fill_count = self._order_event_stats_collector.total_fill_count
    total_fill_qty = self._order_event_stats_collector.total_fill_qty
    fill_ratio_Q = total_fill_qty / total_post_qty
    assert fill_ratio_Q >= 0 and fill_ratio_Q <= 1
    fill_ratio_Q = '%.2f %%' % (fill_ratio_Q * 100)
    fill_ratio_M = total_fill_count / total_post_count
    assert fill_ratio_M >= 0 and fill_ratio_M <= 1
    fill_ratio_M = '%.2f %%' % (fill_ratio_M * 100)
    submit_accept_stats, submit_fill_stats = self._order_event_stats_collector.calc_delay_stats()
    fill_ratio_dict = {
        'immediate_fill_count': submit_fill_stats.immediate_fill_count,
        'accept_then_fill_count': submit_fill_stats.accept_then_fill_count,
        'post_count': total_post_count,
        'post_qty': total_post_qty,
        'fill_count': total_fill_count,
        'fill_qty': total_fill_qty,
        'fill_ratio_M': fill_ratio_M,
        'fill_ratio_Q': fill_ratio_Q
    }
    return fill_ratio_dict

  def order_stats_result(self):
    submit_count = self._order_event_stats_collector.total_submit_count
    reject_count = self._order_event_stats_collector.total_submit_reject_count
    accept_count = self._order_event_stats_collector.total_accept_count
    post_count = self._order_event_stats_collector.total_submit_count - self._order_event_stats_collector.total_submit_reject_count
    ignore_count = post_count - accept_count

    ignore_ratio = ignore_count / post_count
    assert ignore_ratio >= 0 and ignore_ratio <= 1
    ignore_ratio = '%.2f %%' % (ignore_ratio * 100)

    submit_accept_stats, submit_fill_stats = self._order_event_stats_collector.calc_delay_stats()
    submit_accept_delay = submit_accept_stats.submit2accept_delay_total / submit_accept_stats.accept_count / 1e9
    fill_count = submit_fill_stats.accept_then_fill_count + submit_fill_stats.immediate_fill_count
    if fill_count > 0:
      submit_fill_delay = submit_fill_stats.submit2fill_delay_total / fill_count / 1e9
    else:
      submit_fill_delay = float('nan')
    order_stats_dict = {
        'submit_count': submit_count,
        'reject_count': reject_count,
        'accept_count': accept_count,
        'ignore_count': ignore_count,
        'ignore_ratio': ignore_ratio,
        'submit2accept_delay_sec': submit_accept_delay,
        'submit2fill_delay_sec': submit_fill_delay
    }
    return order_stats_dict


def launch(exchange, product, trading_dates, root_dir, machine):
  pnl_data = []
  fill_ratio_data = []
  order_stats_data = []
  for trading_date in trading_dates:
    strategy = AnalyzeStrategy(exchange)
    run_from_og_log_archive(strategy.on_log,
                            'Futures',
                            exchange,
                            trading_date,
                            trading_date + datetime.timedelta(hours=24),
                            root_dir,
                            machine,
                            owner=None)
    pnl_data.append(strategy.pnl_result())
    fill_ratio_data.append(strategy.fill_ratio_result())
    order_stats_data.append(strategy.order_stats_result())
  gen_stats_report(trading_dates, exchange, product, pnl_data, fill_ratio_data, order_stats_data)
  return 0


def aggf(x):
  return ','.join(x.tolist())


def _gen_stats_sub_report(trading_dates, exchange, product, dicts, key_columns, columns):
  datas = {}
  for key in key_columns:
    datas[key] = []
  for idx, trading_date in enumerate(trading_dates):
    datas['date'].append(trading_date.strftime('%Y%m%d'))
    datas['exchange'].append(exchange)
    datas['symbol'].append(product)
  for column in columns:
    datas[column] = []
  for idx, trading_date in enumerate(trading_dates):
    for column in columns:
      datas[column].append(dicts[idx][column])
  df = pandas.DataFrame(datas)
  df = df.round(6)
  for column in columns:
    df[column] = df[column].apply(lambda x: str(x))
  pivot = pandas.pivot_table(df, index=key_columns, values=columns, aggfunc=aggf)
  pivot = pivot[columns]
  return pivot


def gen_stats_report(trading_dates,
                     exchange,
                     symbol,
                     pnl_dicts,
                     fill_ratio_dicts,
                     order_stats_dicts):
  key_columns = ['date', 'exchange', 'symbol']
  pnl_columns = ['realized_pnl', 'latest_position', 'latest_position_price']
  fill_ratio_columns = [
      'immediate_fill_count',
      'accept_then_fill_count',
      'post_count',
      'post_qty',
      'fill_count',
      'fill_qty',
      'fill_ratio_M',
      'fill_ratio_Q'
  ]
  order_stats_columns = [
      'submit_count',
      'reject_count',
      'accept_count',
      'ignore_count',
      'ignore_ratio',
      'submit2accept_delay_sec',
      'submit2fill_delay_sec'
  ]

  pandas.set_option('display.width', 1000)

  print('Pnl stats:')
  print(_gen_stats_sub_report(trading_dates, exchange, symbol, pnl_dicts, key_columns, pnl_columns))

  print('Fill ratio stats:')
  print(
      _gen_stats_sub_report(trading_dates,
                            exchange,
                            symbol,
                            fill_ratio_dicts,
                            key_columns,
                            fill_ratio_columns))

  print('\nOrder event stats:')
  print(
      _gen_stats_sub_report(trading_dates,
                            exchange,
                            symbol,
                            order_stats_dicts,
                            key_columns,
                            order_stats_columns))


def main(_):
  exchange = FLAGS.exchange
  product = FLAGS.product
  trading_date = FLAGS.trading_date
  root_dir = os.path.expanduser(FLAGS.root_dir)
  machine = FLAGS.machine
  assert exchange, '--exchange must be specified.'
  assert product, '--product must be specified.'
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  print('Running for %s %s %s %s ...' % (trading_date, machine, exchange, product))
  sys.stdout.flush()
  trading_dates = resolve_trading_date_range_string(trading_date)
  launch(exchange, product, trading_dates, root_dir, machine)


if __name__ == '__main__':
  flags.DEFINE_string('exchange', None, 'Exchange name.')

  flags.DEFINE_string('product', None, 'Product name.')

  flags.DEFINE_string('trading_date',
                      None,
                      'Trading date in format of YYYYMMDD or YYYYMMDD-YYYYMMDD.')

  flags.DEFINE_string('root_dir', '~/data/og_proto_log', 'root_dir.')

  flags.DEFINE_string('machine', None, 'Instance machine name.')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
